import sqlite3
from config import DATABASE

def get_db_connection():
    """获取数据库连接"""
    conn = sqlite3.connect(DATABASE)
    conn.row_factory = sqlite3.Row
    return conn

def create_annotation_replies_table():
    """创建批注回复表"""
    try:
        conn = get_db_connection()
        
        # 检查表是否已存在
        cursor = conn.execute("""
            SELECT name FROM sqlite_master 
            WHERE type='table' AND name='annotation_replies'
        """)
        
        if cursor.fetchone():
            print("批注回复表已存在")
            conn.close()
            return True
        
        # 创建批注回复表
        conn.execute('''
            CREATE TABLE IF NOT EXISTS annotation_replies (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                annotation_id INTEGER NOT NULL,
                user_id INTEGER NOT NULL,
                content TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (annotation_id) REFERENCES annotations (id),
                FOREIGN KEY (user_id) REFERENCES users (id)
            )
        ''')
        
        # 创建索引
        conn.execute('''
            CREATE INDEX IF NOT EXISTS idx_annotation_replies_annotation 
            ON annotation_replies(annotation_id)
        ''')
        
        conn.execute('''
            CREATE INDEX IF NOT EXISTS idx_annotation_replies_user 
            ON annotation_replies(user_id)
        ''')
        
        # 更新数据库版本
        try:
            current_version = conn.execute("SELECT MAX(version) as version FROM db_version").fetchone()['version']
            new_version = current_version + 1
            
            conn.execute("""
                INSERT INTO db_version (version, description) 
                VALUES (?, ?)
            """, (new_version, "添加批注回复功能"))
            
            print(f"数据库版本已更新到 v{new_version}")
        except Exception as e:
            print(f"更新数据库版本失败: {e}")
        
        conn.commit()
        conn.close()
        
        print("批注回复表创建成功")
        return True
        
    except Exception as e:
        print(f"创建批注回复表失败: {e}")
        return False

if __name__ == '__main__':
    print("批注回复功能数据库迁移工具")
    
    if create_annotation_replies_table():
        print("数据库迁移成功")
    else:
        print("数据库迁移失败")